# -*- coding: utf-8 -*-
"""
Created on Tue Sep 29 20:24:13 2020

@author: DELL
"""
import tensorflow as tf
from tensorflow.keras import layers

class Highway_layers(layers.Layer):
    '''
    构造ELMO模型里面的高速公路层
    filters': [
            [1, 32],
            [2, 32],
            [3, 64],
            [4, 128],
            [5, 256],
            [6, 512],
            [7, 512]
        ]
    
    '''
    def __init__(self,n_filters):
        super().__init__(self)
        self.carrygate_dense=layers.Dense(n_filters,activation='sigmoid')
        self.transform_gate_dense=layers.Dense(n_filters, activation='relu')
        
    def call(self,inputs):
        '''
        我看网上是这个是残差连接的一般形式，但是却没有残差连接有效
        '''
        carrygate=self.carrygate_dense(inputs)
        transformgate=self.transform_gate_dense(inputs)
        
        return carrygate*transformgate+(1.0-carrygate)*inputs
    
#下面是投影层
class ProjectionLayer(layers.Layer):
    '''
    将数据输出为LSTM要求的大小，最终是[batch_size,time_steps,dim=512]
    '''
    
    def __init__(self,lstm_dim=512):
        super().__init__(self)
        self.dense=layers.Dense(lstm_dim,activation='relu')
        
    def call(self,inputs):
        return self.dense(inputs)
    
    
class All_Con_MP_Layers(layers.Layer):
    '''
    该类主要用来做卷积和最大池化操作，并且将七个卷积层经过池化后的输出在最后一个维度拼接起来，最终的输出的大小是
    [batchsize,time_steps,32+32+64+128+256+512+512]的矩阵，然后经过高速公路层和投影层，将矩阵的大小调整为LSTM的
    需求的大小，其实也就是为每个单词生成了一个维度为512的嵌入表示，不过这个嵌入表示是上下文无关的，然后输入给双向LSTM，
    生成上下文相关的词向量
    '''
    
    def __init__(self,filters):
        super().__init__(self)
        
        self.ConvLayers=[layers.Conv2D(num,kernel_size=[1,width]) for i,(width, num) in enumerate(filters)]
        self.MaxPoolLayers=[layers.MaxPool2D(pool_size=(1,50-width+1),strides=(1, 1), padding='valid') for i,(width,num) in enumerate(filters)]
        
    def call(self,inputs):
        conout=[conlayer(inputs) for conlayer in self.ConvLayers]
        mpout=[]
        for i in range(len(conout)):
            mpout.append(tf.squeeze(self.MaxPoolLayers[i](conout[i]),axis=2))#使用maxpooling作用并且在第三个维度也就是axis=2压缩张量，经过池化之后的第二个维度的大小是1
            
        #下面在axis=2粘接张量
        out=mpout[0]
        for i in range(1,len(mpout)):
            out=tf.concat([out,mpout[i]], axis=2)
        return out


class LSTM_Layers(layers.Layer):
    '''
    该类主要用来实现双向LSTM层，并且定义三个参数来将不同的LSTM层输出的隐向量结合起来
    论文中的是直接定义了一个维度为3的隐含层权值，我觉得这样是不合理的，我认为应该是权值应该是随
    不同的句子而发生变化的，因而我这里这定义了一个Dense layer，激活函数使用softmax来输出一个[batch_size,time_steps,3]
    这样做的目的就是输出的权值可以根据不同的句子发生变化。
    '''
    def __init__(self,dim,drop_rate,vocab_size):
        super().__init__(self)
        #下面定义所需要的LSTM层
        self.Lstm_fw_layers1=layers.LSTM(dim,return_sequences=True,go_backwards= False, dropout = drop_rate)
        self.Lstm_bw_layers1=layers.LSTM(dim,return_sequences=True,go_backwards= True, dropout = drop_rate)
        self.Lstm_fw_layers2=layers.LSTM(dim,return_sequences=True,go_backwards= False, dropout = drop_rate)
        self.Lstm_bw_layers2=layers.LSTM(dim,return_sequences=True,go_backwards= True, dropout = drop_rate)
        self.layers_weights=layers.Dense(3, activation='softmax')
        self.outlayer=layers.Dense(vocab_size+1,activation='softmax')
        
    def call(self,inputs):
        self.bilstm1=layers.Bidirectional(merge_mode = "sum", layer =self.Lstm_fw_layers1, backward_layer =self.Lstm_bw_layers1)
        self.bilstm2=layers.Bidirectional(merge_mode = "sum", layer =self.Lstm_fw_layers2, backward_layer =self.Lstm_bw_layers2)
        
        h1=self.bilstm1(inputs)
        h2=self.bilstm2(h1)
        
        #下面计算权重,在这里我选择了将两个隐层和一个输入inputs相加在输入进dense层来计算各层每个隐层和输入的权重
        w=self.layers_weights(inputs+h1+h2)
        w=tf.expand_dims(w, axis=2)
        out=tf.concat([tf.expand_dims(inputs, axis=2),tf.expand_dims(h1, axis=2),tf.expand_dims(h2, axis=2)],axis=2)
        out=tf.squeeze(tf.matmul(w,out),axis=2) 
        
        out=self.outlayer(out)
        
        return out